import requests
from requests_toolbelt import MultipartEncoder
import os
import contextlib
import mimetypes
from dotenv import load_dotenv

load_dotenv()

API_URL = "https://ai.gitee.com/v1/images/ocr"
api_key = os.getenv("api_key")
headers = {
	"Authorization": f"Bearer {api_key}" # 使用 Bearer 令牌进行身份验证 
}

def query(payload):
	fields = [
		("model", payload["model"]),
		("response_format", payload["response_format"]),
	]
	with contextlib.ExitStack() as stack:
		filepath = payload["image"]
		name = os.path.basename(filepath)
		if filepath.startswith(("http://", "https://")):
			response = requests.get(filepath, timeout=10)
			response.raise_for_status()
			fields.append(("image", (name, response.content, response.headers.get("Content-Type", "application/octet-stream"))))
		else:
			mime_type, _ = mimetypes.guess_type(filepath)
			fields.append(("image", (name, stack.enter_context(open(filepath, "rb")), mime_type or "application/octet-stream")))
		encoder = MultipartEncoder(fields)
		headers["Content-Type"] = encoder.content_type
		response = requests.post(API_URL, headers=headers, data=encoder)
		return response.json()

output = query({
	"image": "path/to/image.webp", # 替换为实际的图片路径
	"model": "GOT-OCR2_0",
	"response_format": "text"
})
print(output)